import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.cluster import KMeans

# UserWarning: KMeans is known to have a memory leak on Windows with MKL, when there are less chunks than available threads. You can avoid it by setting the environment variable OMP_NUM_THREADS=1.
OMP_NUM_THREADS = 1

print("DAT Analysis Tool V3.0 by Sheng.")

while input("\nPress Enter to continue, Q to quit...  ") == "":
    try:
        f_dat = input("DAT Filepath: ")
        n_clu = int(input("Cluster Nums: "))

        raw_data = pd.read_csv(
            f_dat, header=None, names=["Name", "Attr", "E", "N", "U"], encoding="utf-8"
        )
        data_set = raw_data[["E", "N", "U"]]
        kmeans = KMeans(n_clusters=n_clu, random_state=0)
        kmeans.fit(data_set)
        labels = kmeans.labels_
        centroids = kmeans.cluster_centers_

        h_stds = []
        v_stds = []
        v_rngs = []
        for i in range(n_clu):
            points = data_set[labels == i]
            stde = np.std(points["E"], axis=0, ddof=1)
            stdn = np.std(points["N"], axis=0, ddof=1)
            stdh = np.std(points["U"], axis=0, ddof=1)
            rngh = max(points["U"]) - min(points["U"])
            h_stds.append(np.sqrt(stdn**2 + stde**2))
            v_stds.append(stdh)
            v_rngs.append(rngh)

        f_rpo = f_dat.replace(".dat", ".md")
        with open(f_rpo, "w", encoding="utf-8") as f:
            f.write(
                f"""## Settings
---

**DAT Filepath:** {f_dat}

**Cluster Nums:** {n_clu}

**Cluster Method:** KMeans

**Plot:** True
"""
            )
            f.write(
                """
## Cluster Labels
---

| 01 | 02 | 03 | 04 | 05 | 06 | 07 | 08 | 09 | 10 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
"""
            )
            for i, label in enumerate(labels, 1):
                f.write(f"| {label:2d} ")
                if i % 10 == 0 or i == len(labels):
                    f.write(f"|\n")

            f.write(
                """
## Cluster Centroids
---

| Label | E | N | U |
| --- | --- | --- | --- |
"""
            )
            for i, c in enumerate(centroids):
                f.write(f"| {i:2d} | {c[0]:.3f} | {c[1]:.3f} | {c[2]:.3f} |\n")

            f.write(
                """
## Horizontal STDs
---

| Label | H-STD |
| --- | --- |
"""
            )
            for i, h_std in enumerate(h_stds):
                f.write(f"| {i:2d} | {h_std:.6f} |\n")

            f.write(
                """
## Vertical STDs
---

| Label | V-STD |
| --- | --- |
"""
            )
            for i, v_std in enumerate(v_stds):
                f.write(f"| {i:2d} | {v_std:.6f} |\n")

            f.write(
                """
## Vertical Ranges
---

| Label | V-Range |
| --- | --- |
"""
            )
            for i, v_rng in enumerate(v_rngs):
                f.write(f"| {i:2d} | {v_rng:.6f} |\n")

        print(f"Done. Report saved to {f_rpo}.")

        plt.xlabel("E")
        plt.ylabel("N", rotation=0, loc='top')
        plt.scatter(data_set["E"], data_set["N"], c=labels, cmap="cool")
        plt.scatter(
            centroids[:, 0],
            centroids[:, 1],
            c=range(n_clu),
            cmap="cool",
            s=100,
            marker="x",
        )
        for i in range(n_clu):
            plt.text(
                centroids[i, 0],
                centroids[i, 1],
                i,
                ha="center",
                va="center",
                fontsize=10,
            )
        plt.grid(True, which="both")
        plt.tight_layout()
        plt.show()

    except Exception as e:
        print(f"Error: {e}")
